import os
import json
from itertools import chain

import numpy as np
import cv2 as cv

from .json_helper import json_helper

class dataset_tools:
    """
    tools to parse dataset sample file and get image
    """

    config = "../../config/dataset_path.json"
    key = "Sentinel12"
    dset_folder = None
    pair_list = None


    def __init__(
        self, dset_dir, list_file
    ) -> None:

        self.config = os.path.normpath(os.path.join(
            __file__, self.config
        ))
        if dset_dir != None and dset_dir != "":
            self.dset_folder = dset_dir
        else:
            jh = json_helper(json_dir=self.config)
            dict_ = jh.get_dict()
            self.dset_folder = dict_[self.key]
        
        self.pair_list = self._parse_list_file(list_file)


    def __len__(self):
        return len(self.pair_list)


    def get_pair_list(self) -> list:
        return self.pair_list


    def get_wkt(self, sample) -> str:
        """
        image name to wkt
        
        Param
        -----
        sample the identifier of sample, e.g. path or point
        """
        return ""


    def get_patch(self, sample, transform=None, **kwargs):
        '''
        sample info, transfrome -> img
        '''

        try:
            img = cv.imread(sample)
        except:
            return None

        if len(img.shape) == 2:
            img = np.expand_dims(img, axis=2)

        if "modal" in kwargs:
            h, w = img.shape[0], img.shape[1]
            if kwargs["modal"] == "SAR":
                img = img[:, :, 0].reshape(h, w, -1)
            if kwargs["modal"] == "OPT":
                img = img.mean(axis=2, keepdims=True).reshape(h, w, -1)

        return transform(img) if transform is not None else img


    def _parse_list_file(self, list_file):
        """
        list file -> dataset dict and list of pairs
        """
        jh = json_helper(json_dir=list_file)
        dict_ = jh.get_dict()

        wkt_dict = {}
        # here is opt before sar, so that elem of result
        # is [OPT_path, SAR_path]
        for f in chain(dict_["OPT"], dict_["SAR"]):
            wkt = self.get_wkt(f)
            if wkt not in wkt_dict:
                wkt_dict[wkt] = []
            wkt_dict[wkt].append(
                os.path.join(self.dset_folder, f)
            )
        
        res = [wkt_dict[k] for k in filter(
                lambda x : len(wkt_dict[x])>1, 
                wkt_dict.keys())]

        return res


class pm_dst_tools:
    
    config = "../../config/dataset_path.json"
    key = "OS"
    dset_folder = None
    pair_list = None

    def __init__(self, dset_dir, list_file) -> None:

        self.config = os.path.normpath(os.path.join(
            __file__, self.config
        ))
        if dset_dir != None and dset_dir != "":
            self.dset_folder = dset_dir
        else:
            jh = json_helper(json_dir=self.config)
            dict_ = jh.get_dict()
            self.dset_folder = dict_[self.key]
        
        self.pair_list = self._parse_list_file(list_file)

    def __len__(self):
        return len(self.pair_list)

    def get_pair_list(self) -> list:
        return self.pair_list

    def _parse_list_file(self, list_file):
        """
        list file -> dataset dict and list of pairs
        """

        with open(list_file) as fp:
            content = json.load(fp)
        
        if not "LABEL" in content:
            length = len(content["OPT"])
            content["LABEL"] = [1] * length

        res = [[
            os.path.join(self.dset_folder, o), 
            os.path.join(self.dset_folder, s), 
            l] for o, s, l in zip(
            content["OPT"], content["SAR"], content["LABEL"]
        )]

        return res

    def get_patch(self, sample, transform=None, **kwargs):
        '''
        sample info, transfrome -> img
        '''

        try:
            img = cv.imread(sample)
        except:
            return None

        if len(img.shape) == 2:
            img = np.expand_dims(img, axis=2)

        if "modal" in kwargs:
            h, w = img.shape[0], img.shape[1]
            if kwargs["modal"] == "SAR":
                img = img[:, :, 0].reshape(h, w, -1)
            if kwargs["modal"] == "OPT":
                img = img.mean(axis=2, keepdims=True).reshape(h, w, -1)

        return transform(img) if transform is not None else img

    # no usage
    def get_wkt(no_usage):
        return 0


class pm_mask_tools(pm_dst_tools):

    def _parse_list_file(self, list_file):
        """
        list file -> dataset dict and list of pairs
        """

        with open(list_file) as fp:
            content = json.load(fp)
        
        if not "LABEL" in content:
            length = len(content["OPT"])
            content["LABEL"] = [1] * length

        res = [[
            os.path.join(self.dset_folder, o), 
            os.path.join(self.dset_folder, s), 
            os.path.join(self.dset_folder, m), 
            l] for o, s, m, l in zip(
            content["OPT"], content["SAR"], content["MASK"], content["LABEL"]
        )]

        return res

    def get_patch(self, sample, transform=None, **kwargs):
        '''
        sample info, transfrome -> img
        '''

        try:
            img = cv.imread(sample)
        except:
            return None

        if len(img.shape) == 2:
            img = np.expand_dims(img, axis=2)

        if "modal" in kwargs:
            h, w = img.shape[0], img.shape[1]
            if kwargs["modal"] == "SAR":
                img = img[:, :, 0].reshape(h, w, -1)
            if kwargs["modal"] == "OPT":
                img = img.mean(axis=2, keepdims=True).reshape(h, w, -1)
            if kwargs["modal"] == "MASK":
                transform = None
                img = img.mean(axis=2, keepdims=True).reshape(h, w, -1)
                img = img / 255

        return transform(img) if transform is not None else img

if __name__ == "__main__1":

    dt = dataset_tools("", "E:/workspace/SOMatch/tmp/json/sen12_tt_harris/pt_s100.json")
    print(f"{len(dt)} pairs loaded")
    lt = dt.get_pair_list()
    img0 = dt.get_patch(lt[0][0])
    img1 = dt.get_patch(lt[0][1])
    img = np.concatenate((img0, img1), axis=1)
    cv.imshow("", img)
    cv.waitKey(100)

if __name__ == "__main__1":

    dt = pm_dst_tools("", "E:/workspace/SOMatch/json/os-select/os-pm-val.json")
    print(f"{len(dt)} pairs loaded")
    lt = dt.get_pair_list()
    print(lt[0][0])
    img0 = dt.get_patch(lt[0][0])
    img1 = dt.get_patch(lt[0][1])
    l = lt[0][2]
    img = np.concatenate((img0, img1), axis=1)
    cv.imshow("", img)
    print(l)
    cv.waitKey(1000)

if __name__ == "__main__":

    dt = pm_mask_tools("", "E:/workspace/SOMatch/json/os-select/1-s-test-mask.json")
    print(f"{len(dt)} pairs loaded")
    lt = dt.get_pair_list()
    print(lt[0][0])
    img0 = dt.get_patch(lt[0][0], modal="OPT")
    img1 = dt.get_patch(lt[0][1], modal="SAR")
    img2 = dt.get_patch(lt[0][2], modal="MASK")
    print(img2.max(), img2.min(), img2.shape)
    # img = np.concatenate((img0*255, img1*255, img2*255), axis=1)
    cv.imshow("", img)
    # print(l)
    cv.waitKey(-1)
